#include <dlib/image_processing/frontal_face_detector.h>  // 人脸检测器
#include <dlib/image_processing.h>  // 图像处理工具
#include <dlib/image_io.h>  // 图像读写
#include <dlib/dnn.h>  // 深度学习模块
#include <iostream>  // 标准输入输出
#include <fstream>  // 文件操作
#include <dirent.h>  // 目录操作
#include <vector>  // 向量容器
#include <algorithm>  // 算法函数
#include <resnet.h>  // ResNet网络定义

using namespace dlib;  // 使用dlib命名空间
using namespace std;  // 使用标准命名空间

anet_type net;  // 定义ResNet人脸识别网络实例

/**
 * 获取指定目录下的所有图片文件
 * @param dir 目录路径
 * @return 包含所有图片文件名的vector
 */
std::vector<string> get_image_files(const string& dir) {
    std::vector<string> files;  // 存储文件名的容器
    DIR *dp;  // 目录指针
    struct dirent *dirp;  // 目录条目
    
    // 尝试打开目录
    if((dp = opendir(dir.c_str())) == NULL) {
        cerr << "无法打开目录: " << dir << endl;
        return files;  // 返回空vector
    }

    // 遍历目录中的文件
    while ((dirp = readdir(dp)) != NULL) {
        string filename = dirp->d_name;
        // 跳过当前目录和上级目录
        if (filename == "." || filename == "..") continue;
        
        // 获取文件扩展名并转为小写
        string ext = filename.substr(filename.find_last_of(".") + 1);
        transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
        
        // 检查是否为支持的图片格式
        if (ext == "jpg" || ext == "jpeg" || ext == "png") {
            files.push_back(filename);  // 添加到结果列表
        }
    }
    closedir(dp);  // 关闭目录
    return files;  // 返回找到的文件列表
}

/**
 * 主函数：构建人脸特征数据库
 */
int main() {
    // 初始化人脸检测器
    frontal_face_detector detector = get_frontal_face_detector();
    
    // 加载人脸关键点检测模型
    shape_predictor sp;
    deserialize("/home/hhy212306170/practice/model/shape_predictor_5_face_landmarks.dat") >> sp;
    
    // 加载人脸识别模型
    deserialize("/home/hhy212306170/practice/model/dlib_face_recognition_resnet_model_v1.dat") >> net;

    // 设置人脸图片目录路径
    string face_dir = "/home/hhy212306170/practice/dlib/myapp/face_images";
    
    // 获取目录下所有图片文件
    std::vector<string> image_files = get_image_files(face_dir);
    if (image_files.empty()) {
        cerr << "目录中没有找到图片文件" << endl;
        return -1;  // 没有图片时退出
    }

    // 创建并打开CSV文件
    ofstream csv_file("facelib.csv");
    if (!csv_file.is_open()) {
        cerr << "无法创建facelib.csv文件" << endl;
        return -1;  // 文件创建失败时退出
    }

    // 处理每张图片
    for (const auto& filename : image_files) {
        // 从文件名提取人名（去掉扩展名）
        string person_name = filename.substr(0, filename.find_last_of("."));
        
        // 构建完整图片路径
        string image_path = face_dir + "/" + filename;
        
        // 加载图片
        matrix<rgb_pixel> img;
        load_image(img, image_path);

        // 检测图片中的人脸
        std::vector<rectangle> faces = detector(img);
        if (faces.empty()) {
            cerr << "图片 " << filename << " 中未检测到人脸" << endl;
            continue;  // 跳过没有检测到脸的图片
        }

        // 检测第一张人脸的关键点
        auto shape = sp(img, faces[0]);
        
        // 对齐人脸（150x150大小，0.25的边界比例）
        matrix<rgb_pixel> face_chip;
        extract_image_chip(img, get_face_chip_details(shape, 150, 0.25), face_chip);

        // 提取128维人脸特征向量
        matrix<float,0,1> face_descriptor = net(face_chip);

        // 将人名和特征向量写入CSV文件
        csv_file << person_name;  // 先写入人名
        for (long i = 0; i < face_descriptor.size(); ++i) {
            csv_file << "," << face_descriptor(i);  // 写入每个特征值
        }
        csv_file << endl;  // 换行
        
        // 打印处理进度
        cout << "已处理: " << person_name << endl;
    }

    // 关闭CSV文件
    csv_file.close();
    cout << "人脸库构建完成，已保存到 facelib.csv" << endl;

    return 0;  // 程序正常结束
}